
#ifndef __CAMERA_BACKBONE_HPP__
#define __CAMERA_BACKBONE_HPP__

#include <memory>
#include <string>
#include <vector>

#include "common/dtype.hpp"

namespace dsgn {

class Head {
 public:
    virtual void forward(const nvtype::half* spatial_feature_2d,void* stream = nullptr) = 0;

    virtual nvtype::half* box_preds() = 0;
    virtual nvtype::half* cls_preds() = 0;
    virtual nvtype::half* dir_cls_preds() = 0;
    virtual std::vector<int> box_preds_shape() = 0;
    virtual std::vector<int> cls_preds_shape() = 0;
    virtual std::vector<int> dir_cls_preds_shape() = 0;
    virtual void print() = 0;
};

std::shared_ptr<Head> create_dethead(const std::string& model);

};  // namespace dsgn

#endif  // __CAMERA_BACKBONE_HPP__